import cv2
import numpy as np


def extend_border(image, top, bottom, left, right, border_type=cv2.BORDER_CONSTANT, value=0):
    return cv2.copyMakeBorder(image, top, bottom, left, right, border_type, value=value)


def apply_optical_flow_with_border(image_path, flow_path, output_image_path, border_size=256):
    try:
        # 加载图像
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"无法打开图像文件：{image_path}")

        # 加载光流场
        flow = np.load(flow_path)

        # 扩充图像边界
        image_with_border = extend_border(image, border_size, border_size, border_size, border_size,
                                          cv2.BORDER_CONSTANT, value=[0, 0, 0])

        # 生成变形后的坐标
        h_b, w_b = image_with_border.shape[:2]
        map_x, map_y = np.meshgrid(np.arange(w_b), np.arange(h_b))
        map_x = map_x.astype(np.float32) + flow[:, :, 0]
        map_y = map_y.astype(np.float32) + flow[:, :, 1]

        # 应用光流进行变形
        remapped_image = cv2.remap(image_with_border, map_x, map_y, interpolation=cv2.INTER_LINEAR,
                                   borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))

        # 保存结果图像
        cv2.imwrite(output_image_path, remapped_image)

        print(f"Processed image saved to {output_image_path}")

    except FileNotFoundError as e:
        print(f"文件未找到: {e}")
    except PermissionError as e:
        print(f"权限错误: {e}")
    except Exception as e:
        print(f"发生错误: {e}")


# 示例用法
image_path = r'E:\wafer52\11867_256nm\S11867M1-80_tr1-tc1.png'  # 替换为实际路径
flow_path = r'E:\wafer52\256nm_flow_combine\large_flow_repair.npy'  # 替换为实际路径
output_image_path = r'E:\wafer52\11866_256nm\S11867M1-80_tr1-tc1_repair.png'  # 替换为实际路径

apply_optical_flow_with_border(image_path, flow_path, output_image_path,border_size=0)

print(f"Processed image saved to {output_image_path}")
